"""
数据库数据加载器
"""
import pandas as pd
from typing import Dict, Any
from .base_loader import BaseLoader


class DBLoader(BaseLoader):
    """
    数据库数据加载器
    """
    def __init__(self, db_connection: str):
        self.db_connection = db_connection
    
    def load(self, source: str, **kwargs) -> Dict[str, Any]:
        """
        加载数据库数据
        
        Args:
            source: SQL 查询语句
            **kwargs: 查询参数
        """
        data = pd.read_sql(source, self.db_connection, **kwargs)
        return data.to_dict(orient="records")